0.8.0 — a state holds only itself, and a step can give it up - #29

Released as 0.8.0 together with #28. The version bump is the last commit here. Merge #28 first, then this; the two combine without conflicts, and the combined tree passes 650 tests.

Closes #26.

By default, SparseDeltaMemory.step returns a new table per stream: O(B·N·d_v) a step, against O(B·(W+R)·d_v) of arithmetic. At a large batch that copy is most of a decode step, even at ordinary table sizes. This PR adds a way to waive it:

y, state = trunk.step(x, state, donate=True)   # the old state is spent

donate=True says the caller will not read state again, so a sub-layer may write its successor into the buffers it was given. The default is unchanged: every caller who doesn't opt in keeps the successor guarantee.

Shape

• A keyword, not a step_ method. The path that needs it is a trunk stepping many streams, so the option has to pass through Block.step and Stack.step. A keyword does that, and every component keeps the same three methods.
• Gated DeltaNet and Undertow accept the keyword and ignore it. Their successors cost no more than their own arithmetic.
• Block passes the keyword on only when it is set. A stateful sub-layer from outside the library that doesn't take it keeps stepping, and fails loudly only if someone asks it to donate.
• Not copy-on-write paging. Product keys spread writes across the table on purpose, so with 64-row pages, 64 writes touch essentially every page of a 32² table.

Donating is a permission, not a demand

recurrent_sparse_delta takes in_place. Both paths run one body whose last write is either scatter or scatter_, so they agree bit for bit. The layer still copies when the table is not the stream's own to write:

• Any view. This includes init_state's broadcast. Under "learned" that broadcast shares the parameter's storage, and at B = 1 it does not overlap, so an in-place write would go through and change the learned table.
• A table autograd is recording.
• A table under a torch.func transform.

So the first donated step after init_state is a copy, and every step after it writes in place.

Docs

• Design record §3.6 now carries the measurements and the reasons for this shape. Its §8 open item on an in-place decode is closed.
• Four places said per-token work does not depend on N, with no qualification: the module docstring, the n_slots docstring, the package docstring and docs/SDM.md. They now say the arithmetic does not, and the successor copy does.
• docs/SDM.md said "at small tables it does not matter", which is false at B = 128. Its Generation section and docs/BLOCK.md now document donate.

Measured

One machine (Tesla P40, fp32), through the API: step(x, state) against step(x, state, donate=True). Median of 3 runs of 64 steps each, after 8 warm-up steps. Outputs were bit-equal in every row.

One layer, d_model=512, n_heads=2, W=R=64, initial_memory="learned":

Stack.step with 5 blocks, d_model=512, d_mlp=1024, SDM at blocks 1 and 3 and Gated DeltaNet elsewhere (each stack value is stable across 5 repeats):

Tests

• SDM:
  • a donated stream equals a kept one bit for bit, for both initial_memory options at B = 1 and B = 3;
  • the learned table is untouched after donated steps from init_state, under both no_grad and inference_mode;
  • after the first step the write really is in place, checked by comparing the predecessor's and successor's buffers while both are alive, so a reused allocation can't pass for one;
  • donating is declined under autograd and under vmap.
• Stack: a donated stream equals a kept one bit for bit, and the in-place write reaches the table.
• Block: donate reaches every stateful sub-layer only when set, and a sub-layer without the keyword still steps.

Each guard was checked by removing it: dropping the view check, dropping the transform check, ignoring donate, or passing it unconditionally each fails the test meant to catch it.

• GPU: 645 passed, no skips.
• CPU, CUDA hidden: 579 passed, 66 skipped.
• The first commit passes on its own.
• Merged with state: a state returned by forward holds only itself #28, the tree auto-merges and passes 650.

🤖 Generated with Claude Code